{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE TypeInType #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# OPTIONS_GHC -Wno-redundant-constraints #-}

-- | Module: Data.Vector.Sized.Sample
-- Copyright: (c) 2018 Koz Ross
-- Maintainer: Koz Ross <koz.ross@retro-freedom.nz>
-- Stability: Experimental
-- Portability: Non-portable
--
-- Implementations of various sampling algorithms, both weighted and unweighted.
-- These return new copies.
module Data.Vector.Sized.Sample
  ( algorithmR
  , chao
  ) where

import Control.Monad.Random.Strict
import Control.Monad.ST
import Data.Foldable
import Data.Proxy
import GHC.TypeNats

import qualified Data.Vector.Generic as VG
import qualified Data.Vector.Generic.Mutable.Sized as VGMS
import qualified Data.Vector.Generic.Sized as VGS

-- | Random sampling using Algorithm R.
--
-- The implementation is based on Vitter's paper ["Random sampling with a
-- reservoir"](http://www.cs.umd.edu/~samir/498/vitter.pdf). 
--
-- __Complexity:__ \(O(n)\)
{-# INLINE algorithmR #-}
algorithmR ::
     forall k m v n a l.
     (KnownNat n, KnownNat k, MonadRandom m, VG.Vector v a, (k + l) ~ n)
  => VGS.Vector v n a
  -> m (VGS.Vector v k a)
algorithmR v = do
  ixes <- traverse (\i -> (i, ) <$> getRandomR (0, i)) [k' .. (n' - 1)]
  let v' =
        runST
          (do vec <- VGMS.unsafeNew
              VGS.copy vec (VGS.take @_ @k v)
              traverse_
                (\(i, j) ->
                   when (j < k') (VGMS.unsafeWrite vec j (VGS.unsafeIndex v i)))
                ixes
              VGS.unsafeFreeze vec)
  return v'
  where
    k' = fromIntegral . natVal $ Proxy @k
    n' = VGS.length v

-- | Weighted random sampling using Chao's algorithm. The input vector is
-- considered to be full of item-weight pairs.
--
-- The algorithm is based on the following approach to weighting: let the weight
-- of two items \(i,j\) be \(w_i,w_j\), and \(p_i\) be the probability that
-- \(i\) is selected; then \(p_j = \min\{1, p_i \cdot \frac{w_j}{w_i}\}\). The
-- implementation is based on the description given by Chao in the paper ["A
-- general-purpose unequal probability sampling
-- plan"](http://biomet.oxfordjournals.org/content/69/3/653.abstract).
--
-- __Complexity:__ \(O(n)\)
{-# INLINE chao #-}
chao ::
     forall k m v n a b l.
     ( KnownNat n
     , KnownNat k
     , MonadRandom m
     , Fractional b
     , Ord b
     , Random b
     , VG.Vector v (a, b)
     , VG.Vector v a
     , (k + l) ~ n
     )
  => VGS.Vector v n (a, b)
  -> m (VGS.Vector v k a)
chao v = do
  ixes <-
    traverse
      (\i -> (i, , ) <$> getRandomR (0.0, 1.0) <*> getRandomR (0, k' - 1))
      [k' .. (n' - 1)]
  let v' =
        runST
          (do vec <- VGMS.unsafeNew
              wSum <-
                foldlM
                  (\s i -> do
                     let (e, w) = VGS.unsafeIndex v i
                     VGMS.unsafeWrite vec i e
                     return (s + w / fromIntegral k'))
                  0.0
                  [0 .. (k' - 1)]
              foldM_
                (\s (i, j, k'') -> do
                   let (e, w) = VGS.unsafeIndex v i
                   let s' = s + w / fromIntegral k'
                   let p = w / s'
                   when (j <= p) (VGMS.unsafeWrite vec k'' e)
                   return s')
                wSum
                ixes
              VGS.unsafeFreeze vec)
  return v'
  where
    k' = fromIntegral . natVal $ Proxy @k
    n' = VGS.length v
